from openai import OpenAI
from config.config import OPENAPI_API_KEY_4, OPENAPI_API_BASE_4
import json
import time
from typing import Callable, Any


class SmartAgent:
    def __init__(self):
        self.client = OpenAI(
            api_key=OPENAPI_API_KEY_4,
            base_url=OPENAPI_API_BASE_4
        )
        self.messages = []
        self.tools = []
        self.functionMap = {}

    def registerTool(self, name: str, description: str, parameters: object, function: Callable) -> None:
        self.tools.append({
            "type": "function",
            "function": {
                "name": name,
                "description": description,
                "parameters": parameters,
            }
        })
        self.functionMap[name] = function

    def run(self, prompt: str) -> str:
        self.messages.append({
            "role": "user",
            "content": prompt
        })
        while True:
            response = self.client.chat.completions.create(
                model="gpt-4o-mini",
                messages=self.messages,
                tools=self.tools,
                tool_choice="auto"
            )
            reply = response.choices[0].message
            if reply.tool_calls:
                for tool in reply.tool_calls:
                    args = json.loads(tool.function.arguments)
                    result = self.functionMap[tool.function.name](**args)
                    self.messages.append({
                        "role": "assistant",
                        "tool_calls": [tool]
                    })
                    self.messages.append({
                        "role": "tool",
                        "name": tool.function.name,
                        "tool_call_id": tool.id,
                        "content": result
                    })
                    time.sleep(0.1)
                    print(result)
            else:
                ## 工具结束
                chatResult = reply.content
                print(reply.content)
                break
        return chatResult
